# get default parameters the papermill way.
            input_file = "../results/04_annotate_cell_types/adata.h5ad"
            output_file = "tmp/data.h5ad"
            output_file_obs = "tmp/adata_obs.tsv"
            output_file_norm_counts = "tmp/norm_counts.tsv"
            results_dir = "tmp"
            table_dir = "../tables"

# Parameters
            input_file = "input_adata.h5ad"
            output_file = "adata.h5ad"
            output_file_obs = "adata_obs.tsv"
            output_file_norm_counts = "norm_counts.tsv"
            cpus = "1"
            table_dir = "tables"
            results_dir = "."

import pandas as pd
            import scanpy as sc
            import numpy as np
            from matplotlib import pyplot as plt
            import sys
            import os
            
            sys.path.append("lib")
            sys.path.append("../lib")
            
            from jupytertools import fix_logging, display
            
            from operator import or_
            from functools import reduce
            
            fix_logging(sc.settings)

markers = pd.read_csv(os.path.join(table_dir, "cell_type_markers.csv"))

cell_types = markers["cell_type"].unique()

adata = sc.read_h5ad(input_file)

sc.pl.umap(adata, color="cell_type")

# PCA turned out not to be entirely reproducible on different CPU architechtures.
            # For the sake of reproducibility of these notebooks, we load a pre-computed result
            # from the repository. If it doesn't exist, we compute it from scratch.
            def pca_precomputed(adata, key, **kwargs):
                pca_file = os.path.join(table_dir, f"adata_pca_{key}.pkl.gz")
                try:
                    # CSV does not preserver floating point numbers perfectly. 
                    adata.obsm["X_pca"] = pd.read_pickle(pca_file).values
                except IOError:
                    assert False, "should use pre-computed version. "
                    sc.tl.pca(adata, svd_solver="arpack", **kwargs)
                    pd.DataFrame(adata.obsm["X_pca"]).to_pickle(pca_file)

adata = adata[adata.obs["cell_type"].isin(["T cell", "NK cell"]), :].copy()

prefixes = ["TRAV", "TRBV", "TRBJ", "TRAJ", "TRBD"]
            exclude_genes = reduce(or_, [adata.var_names.str.startswith(x) for x in prefixes])

adata.var_names[exclude_genes]

adata.shape

adata = adata[:, ~exclude_genes].copy()

adata.shape

np.sum(adata.var_names.str.startswith("TRAV"))

random_state = 1
            sc.pp.neighbors(adata, random_state=random_state)
            sc.tl.umap(adata, random_state=random_state)

sc.tl.leiden(adata, resolution=3, random_state=42)

fig, ax = plt.subplots(figsize=(7, 5))
            sc.pl.umap(
                adata, legend_loc="on data", color="leiden", ax=ax, size=20, legend_fontoutline=3
            )

for ct in cell_types:
                if not "T cell" in ct and ct != "NK cell":
                    continue
                marker_genes = markers.loc[markers["cell_type"] == ct, "gene_identifier"]
                sc.pl.umap(
                    adata,
                    color=marker_genes,
                    title=["{}: {}".format(ct, g) for g in marker_genes],
                    size=15,
                )

fig, ax = plt.subplots(figsize=(7, 5))
            sc.pl.umap(
                adata, legend_loc="on data", color="leiden", ax=ax, size=20, legend_fontoutline=3
            )

annotation_t = {
                "T CD4+": [],
                # 27 is NK, but some of the cells are CD8
                "T CD8+": [8, 0, 6, 20, 11, 29, 21, 22, 14, 27],
                "T reg.": [13, 12],
                "T other": [31],
                "ambiguous": [24, 25, 26, 18],
            }

annot_dict = {str(c): ct for ct, clusters in annotation_t.items() for c in clusters}
            adata.obs["cell_type2"] = [
                annot_dict.get(c, "T CD4+") if ct != "NK cell" else ct
                for c, ct in zip(adata.obs["leiden"], adata.obs["cell_type"])
            ]

sc.pl.umap(adata, color="cell_type2")

adata_ambiguous = adata[adata.obs["cell_type2"] == "ambiguous", :].copy()

sc.__version__

pca_precomputed(adata_ambiguous, "subcluster_ambiguous")
            sc.pp.neighbors(adata_ambiguous, n_neighbors=5)
            sc.tl.umap(adata_ambiguous)

sc.tl.leiden(adata_ambiguous, resolution=2)

sc.pl.umap(
                adata_ambiguous, color=["leiden"], legend_loc="on data", legend_fontoutline=4
            )
            sc.pl.umap(adata_ambiguous, color=["CD8A", "CD8B", "FOXP3", "CD4", "n_counts"], ncols=3)

annotation_ambiguous = {
                "T CD4+": [21, 14, 12, 2, 23, 11, 13, 8, 9, 17, 24, 0, 22, 10],
                "T CD8+": [4, 18, 15, 26],
                "T reg.": [6, 1, 5, 19, 25, 3, 16, 7],
                "T other": [20],
            }

annot_dict = {
                str(c): ct for ct, clusters in annotation_ambiguous.items() for c in clusters
            }
            adata_ambiguous.obs["cell_type3"] = [
                annot_dict[c]
                for c, ct in zip(adata_ambiguous.obs["leiden"], adata_ambiguous.obs["cell_type"])
            ]

sc.pl.umap(adata_ambiguous, color="cell_type3")

# convert categorical to string
            adata.obs["cell_type2"] = [str(x) for x in adata.obs["cell_type2"]]
            adata.obs.loc[adata_ambiguous.obs_names, "cell_type2"] = adata_ambiguous.obs[
                "cell_type3"
            ]

adata.obs["cell_type"] = adata.obs["cell_type2"]
            del adata.obs["cell_type2"]

sc.pl.umap(
                adata,
                color=[
                    "samples",
                    "n_genes",
                    "mt_frac",
                    "has_ir",
                    "hpv_status",
                    "ir_status",
                    "chain_pairing",
                    "phase",
                    "leiden",
                    "CD4",
                    "FOXP3",
                    "CD8A",
                    "CD19",
                    "is_doublet",
                    "cell_type",
                ],
                ncols=3,
            )

adatas = dict()
            resolutions = {"T CD4+": 1, "T CD8+": 1, "NK cell": 0.5, "T reg.": 0.5, "T other": 0.5}
            for ct in adata.obs["cell_type"].unique():
                tmp_adata = adata[adata.obs["cell_type"] == ct, :].copy()
                pca_precomputed(
                    tmp_adata,
                    f'subcluster_{ct.lower().replace(" ", "_").replace("+", "").replace(".", "")}',
                )
                sc.pp.neighbors(tmp_adata)
                sc.tl.umap(tmp_adata)
                sc.tl.leiden(tmp_adata, resolution=resolutions[ct])
                sc.pl.umap(tmp_adata, color="leiden", title=ct)
                adatas[ct] = tmp_adata

for ct, tmp_adata in adatas.items():
                adata.obs.loc[tmp_adata.obs.index, "cluster"] = [
                    "{} {}".format(ct, x) for x in tmp_adata.obs["leiden"]
                ]

sc.pl.umap(adata, color=["cluster", "cell_type"])

adata.write_h5ad(output_file, compression="lzf")
            adata.obs.to_csv(output_file_obs, sep="\t")